from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END

# 模拟子图 child_graph 的定义和编译（这里用简化版本）
class ChildState(TypedDict):
    my_child_key: str

def child_node_1(state: ChildState):
    print(f"🔧 子图节点1: 处理输入 '{state['my_child_key']}'")
    return {"my_child_key": f"子图处理: {state['my_child_key']}"}

def child_node_2(state: ChildState):
    print(f"🔧 子图节点2: 进一步处理 '{state['my_child_key']}'")
    return {"my_child_key": f"{state['my_child_key']} -> 完成处理"}

# 构建子图
print("🏗️ 构建子图...")
child_builder = StateGraph(ChildState)
child_builder.add_node("child_1", child_node_1)
child_builder.add_node("child_2", child_node_2)
child_builder.add_edge(START, "child_1")
child_builder.add_edge("child_1", "child_2")
child_builder.add_edge("child_2", END)

child_graph = child_builder.compile()
print("✅ 子图编译完成")

# 定义父图
class ParentState(TypedDict):  # 父图的状态结构体
    my_key: str  # 父图状态键: my_key

def call_child_graph(state: ParentState) -> ParentState:  # 节点函数，输入输出状态类型都为 ParentState
    print(f"🔄 调用子图节点: 接收父图状态 '{state['my_key']}'")

    # 状态转换: 父图状态 -> 子图状态 (ParentState.my_key -> ChildState.my_child_key)
    child_graph_input = {"my_child_key": state["my_key"]}  # 将父图状态的 my_key 值，赋值给子图状态的 my_child_key 键
    print(f"🔀 状态转换: 父图状态 -> 子图状态")
    print(f"   输入子图: {child_graph_input}")

    # 调用子图
    child_graph_output = child_graph.invoke(child_graph_input)  # 调用子图 child_graph
    print(f"📤 子图输出: {child_graph_output}")

    # 状态转换: 子图状态 -> 父图状态 (ChildState.my_child_key -> ParentState.my_key)
    result = {"my_key": child_graph_output["my_child_key"]}  # 将子图输出状态的 my_child_key 值，赋值给父图状态的 my_key 键
    print(f"🔀 状态转换: 子图状态 -> 父图状态")
    print(f"   返回父图: {result}")

    return result

def parent_node_1(state: ParentState):
    print(f"🔧 父图节点1: 预处理 '{state['my_key']}'")
    return {"my_key": f"父图预处理: {state['my_key']}"}

def parent_node_2(state: ParentState):
    print(f"🔧 父图节点2: 后处理 '{state['my_key']}'")
    return {"my_key": f"{state['my_key']} -> 父图完成"}

print("\n🏗️ 构建父图...")
builder = StateGraph(ParentState)  # 创建父图 StateGraph，指定状态结构体为
ParentState

# 父图的其他节点定义
builder.add_node("parent_1", parent_node_1)
builder.add_node("parent_2", parent_node_2)
builder.add_node("child", call_child_graph)  # 将节点函数 call_child_graph 作为节点添加到父图，节点名为 "child"

builder.add_edge(START, "parent_1")  # 父图节点 parent_1 -> 节点函数 call_child_graph (普通边)
builder.add_edge("parent_1", "child")
builder.add_edge("child", "parent_2")
builder.add_edge("parent_2", END)

graph = builder.compile()  # 编译父图
print("✅ 父图构建完成！")

# 测试运行
initial_state = {"my_key": "用户输入数据"}
print(f"\n📥 初始父图状态: {initial_state}")
print("\n" + "="*60)

result = graph.invoke(initial_state)

print("="*60)
print(f"\n📤 最终父图状态: {result}")